[python][torch] Add batch-first streaming datasets - #9365
Conversation
| return True | ||
| total_rows = 0 | ||
| for split in self.splits: | ||
| row_count = getattr(split, "row_count", None) |
There was a problem hiding this comment.
[P2] Use merged row counts before collapsing DataLoader workers
Split.row_count is the sum of physical file rows, so it can substantially exceed the logical output for deletion-vector and Data Evolution splits. For example, if two splits each have 10 physical rows but 4 merged rows and the limit is 8, the limit is non-binding; this code nevertheless totals 20 and routes every split to worker 0, unnecessarily regressing row, shuffle, and batch streaming from multi-worker to single-worker execution. TableRead._limit_covers_all_splits already uses split.merged_row_count() for the same proof. Please prefer a validated merged_row_count() when it is available, fall back conservatively when it is unknown, and add a non-binding test where physical and merged counts differ.
There was a problem hiding this comment.
Fixed in cc695c0. The worker-limit proof now prefers merged_row_count() only when it is a valid nonnegative count no larger than the physical row count; otherwise it falls back to the physical upper bound. The regression covers 20 physical rows, 8 merged rows, and limit=8 retaining both workers.
2 participants
Purpose
Avoid per-row Python conversion by allowing PyPaimon Torch streaming datasets to yield PyArrow
RecordBatchobjects or dictionaries of Torch tensors.batch_format="row"remains the default;"pyarrow"and"torch"enable batch streaming.batch_size=Nonepreserves reader batches; otherwise batches are combined or sliced.to_tensor_fn.DataLoader(batch_size=None)to avoid a second batching step.Batch formats use one reader per DataLoader worker and require
prefetch_concurrency=1, because Python producer threads cannot reliably cancel a reader blocked in I/O. A non-binding global limit preserves DataLoader worker parallelism when validated merged row counts, or physical row-count upper bounds, prove that all rows fit.Related implementations
LanceDataset also streams Arrow batches and provides
batch_size,to_tensor_fn, and nativebatch_readahead. PyIceberg providesto_arrow_batch_reader(), but not a PyTorch Dataset or Tensor conversion layer.Tests
py_compile, andgit diff --check